# scripts/import_data.py
import os
import pandas as pd
from django.conf import settings
import django

# 设置 Django 环境
os.environ.setdefault('DJANGO_SETTINGS_MODULE', 'backend.settings')
django.setup()

from problems.models import Problem, Tag

# 获取文件路径
tag_map_file = os.path.join(settings.BASE_DIR, 'static', 'data', 'TagMap.xlsx')
problem_file_b = os.path.join(settings.BASE_DIR, 'static', 'data', 'ProblemsInfoB.xlsx')
problem_file_p = os.path.join(settings.BASE_DIR, 'static', 'data', 'ProblemsInfoP.xlsx')

# 读取标签映射数据
tags_df = pd.read_excel(tag_map_file, engine='openpyxl')

# 导入标签
for index, row in tags_df.iterrows():
    tag, created = Tag.objects.get_or_create(
        index=row['Index'],
        tag_name=row['Tag']
    )
    if created:
        print(f"Tag {tag.tag_name} created.")


# 定义一个函数来导入题目数据
def import_problems_from_excel(file_path):
    problems_df = pd.read_excel(file_path, engine='openpyxl')

    for index, row in problems_df.iterrows():
        # 处理标签映射，假设标签是逗号分隔的字符串，如 "1,2,3"
        tags_str = row['Tags']
        if isinstance(tags_str, str):
            tag_ids = [int(tag.strip()) for tag in tags_str.split(',') if tag.strip().isdigit()]
        else:
            tag_ids = []  # 如果标签不是字符串，则为空

        # 更新或导入题目信息
        problem, created = Problem.objects.update_or_create(
            problem_number=row['Problem Number'],
            defaults={
                'problem_name': row['Problem Name'],
                'difficulty': row['Difficulty'],
                'submission_count': row['Submission Count'],
                'accepted_count': row['Accepted Count'],
                'acceptance_rate': row['Acceptance Rate'],
            }
        )

        # 添加标签到题目
        for tag_id in tag_ids:
            try:
                tag = Tag.objects.get(index=tag_id)
                problem.tags.add(tag)
            except Tag.DoesNotExist:
                print(f"Tag with index {tag_id} does not exist.")

        problem.save()
        print(f"Problem {problem.problem_number} imported.")


# 导入两个问题文件中的数据
import_problems_from_excel(problem_file_b)  # 导入 ProblemsInfoB.xlsx
import_problems_from_excel(problem_file_p)  # 导入 ProblemsInfoP.xlsx

print("Data import complete.")